//
//  RSReplay.m
//  ReplayStar
//
//  Created by Félix Cloutier on 11-11-02.
//  Copyright 2011 __MyCompanyName__. All rights reserved.
//

#import <ApplicationServices/ApplicationServices.h>

#import "RSReplay.h"
#import "StormLib.h"
#import "RSPlayerData.h"
#import "RSChatMessage.h"
#import "SOIL.h"
#import "stb_image_aug.h"
#import "stbi_DDS_aug.h"

#undef check
#import "checked_ptr.h"

#define NSN(x) [NSNumber numberWithInt:(x)]

typedef struct RSReplayAttributesData
{
	unsigned header; // e7 03
	unsigned attributeId; // xx yy 00 00
	unsigned char playerId;
	unsigned value;
	
} __attribute__((packed)) RSReplayAttributesData;

typedef struct RSReplayAttributesEventHeader
{
	unsigned header; // 00 00 00 00
	unsigned attributesCount;
	RSReplayAttributesData attributes[];
	
} __attribute__((packed)) RSReplayAttributesEventHeader;

static NSString* kRPReplayDetailsKey = @"replay.details";
static NSString* kRPAttributeEventsKey = @"replay.attributes.events";
static NSString* kRPReplayInitDataKey = @"replay.initData";
static NSString* kRPReplayGameEventsKey = @"replay.game.events";
static NSString* kRPReplayMessageEventsKey = @"replay.message.events";

static NSDictionary* fileParsers;

static NSString* kRSMapDataPathFormat = @"/Users/Shared/Blizzard/Battle.net/Cache/%02hhx/%02hhx/%@.s2ma";

static CGImageRef loadDDS(NSData* data)
{
	int width, height, channels;
	unsigned char* imageBytes = SOIL_load_image_from_memory((const unsigned char*)data.bytes, (int)data.length, &width, &height, &channels, SOIL_LOAD_RGBA);
	
	CGColorSpaceRef rgb = CGColorSpaceCreateDeviceRGB();
	CGContextRef bitmapCtx = CGBitmapContextCreate(imageBytes, width, height, 8, width * 4, rgb, kCGImageAlphaPremultipliedLast);
	assert(bitmapCtx != NULL && "Null bitmapCtx");
	
	CGImageRef result = CGBitmapContextCreateImage(bitmapCtx);
	
	CFRelease(rgb);
	CFRelease(bitmapCtx);
	free(imageBytes);
	
	[(id)result autorelease];
	return result;
}

static CGImageRef loadTGA(NSData* data)
{
	NSString* typeHintKey = (NSString*)kCGImageSourceTypeIdentifierHint;
	NSDictionary* typeHint = [NSDictionary dictionaryWithObject:@"com.truevision.tga-image" forKey:typeHintKey];
	CGImageSourceRef targaSource = CGImageSourceCreateWithData((CFDataRef)data, (CFDictionaryRef)typeHint);
	CGImageRef image = CGImageSourceCreateImageAtIndex(targaSource, 0, NULL);
	CFRelease(targaSource);
	[(id)image autorelease];
	return image;
}

static NSString* stringnifyMapHash(NSData* mapHash)
{
	static const char hegits[] = {'0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f'};
	const char* bytes = (const char*)mapHash.bytes;
	char hash[65];
	hash[64] = 0;
	for (int i = 0; i < 32; i++)
	{
		unsigned char byte = bytes[i];
		hash[i * 2] = hegits[byte >> 4];
		hash[i * 2 + 1] = hegits[byte & 0xF];
	}
	
	return [NSString stringWithCString:hash encoding:NSASCIIStringEncoding];
}

static NSDate* dateFromSC2Timestamp(unsigned long long ts)
{
	NSTimeInterval time = (ts - 116444735995904000LL) / 10000000;
	return [NSDate dateWithTimeIntervalSince1970:time];
}

static unsigned parseTimeDelta(checked_ptr<const unsigned char>& data)
{
	unsigned char lengthBits = data[0] & 3;
	unsigned result = data[0] >> 2;
	data++;
	for (int i = 0; i < lengthBits; i++)
	{
		result <<= 8;
		result |= *data;
		data++;
	}
	return result;
}

static long long parseHorribleNumber(checked_ptr<const unsigned char>& bytes)
{
	long long retval = 0;
	BOOL hasNext = YES;
	
	checked_ptr<const unsigned char> iter = bytes;
	for (size_t i = 0; hasNext; i++)
	{
		unsigned char byte = *iter;
		unsigned long long byteValue = byte & 0x7F;
		unsigned long long shift = 7 * i;
		hasNext = byte >> 7;
		
		retval |= byteValue << shift;
		iter++;
	}
	
	BOOL negative = retval & 1;
	retval >>= 1;
	if (negative) retval *= -1;
	
	bytes = iter;
	return retval;
}

static id parseHorribleStructure(checked_ptr<const unsigned char>& bytes)
{
	unsigned char byte = *bytes;
	bytes++;
	switch (byte)
	{
		case 2: // binary
		{
			long long length = parseHorribleNumber(bytes);
			
			bytes.check(length);
			// assume a valid UTF8 sequence; if not, fallback to data
			id result = [[[NSString alloc] initWithBytes:bytes.to_unchecked() length:length encoding:NSUTF8StringEncoding] autorelease];
			if (result == nil)
			{
				result = [NSData dataWithBytes:bytes.to_unchecked() length:length];
			}
			bytes += length;
			return result;
		}
			
		case 4: // array
		{
			assert(bytes[0] == 1 && bytes[1] == 0 && "Unexpected magic values");
			bytes += 2;
			long long length = parseHorribleNumber(bytes);
			NSMutableArray* array = [NSMutableArray array];
			for (long long i = 0; i < length; i++)
			{
				id object = parseHorribleStructure(bytes);
				[array addObject:object];
			}
			return [[array copy] autorelease];
		}
			
		case 5: // map
		{
			long long length = parseHorribleNumber(bytes);
			NSMutableDictionary* dict = [NSMutableDictionary dictionary];
			for (long long i = 0; i < length; i++)
			{
				long long index = parseHorribleNumber(bytes);
				id object = parseHorribleStructure(bytes);
				[dict setObject:object forKey:[NSNumber numberWithLongLong:index]];
			}
			return [[dict copy] autorelease];
		}
			
		case 6: // byte
		{
			unsigned char byte = *bytes;
			bytes++;
			return [NSNumber numberWithUnsignedChar:byte];
		}
			
		case 7: // integer
		{
			const int* ptr = reinterpret_cast<const int*>(bytes.to_unchecked());
			bytes += 4;
			return [NSNumber numberWithInt:*ptr];
		}
			
		case 9: // number
		{
			long long result = parseHorribleNumber(bytes);
			return [NSNumber numberWithLongLong:result];
		}
			
		default:
			assert(!"Unknown data type");
			return nil;
	}
}

@interface RSReplay (Private)

-(id)crippledInitWithContentsOfFile:(NSString*)file;
-(void)loadPreviewFromCacheFiles:(NSArray*)cacheFiles previewName:(NSString *)previewFilePath;
-(NSData*)readArchivedFile:(NSString*)file;

@end

@implementation RSReplay

@synthesize duration;
@synthesize startDate;
@synthesize endDate;
@synthesize mapName;
@synthesize mapPreview;
@synthesize messages;
@synthesize teams;
@synthesize speed;
@synthesize type;
@synthesize ranking;
@synthesize version;

+(CGImageRef)previewOfFile:(NSString *)file
{
	RSReplay* crippled = [[[self alloc] crippledInitWithContentsOfFile:file] autorelease];
	
	if (NSData* detailsData = [crippled readArchivedFile:kRPReplayDetailsKey])
	{
		const unsigned char* rawBytes = static_cast<const unsigned char*>(detailsData.bytes);
		checked_ptr<const unsigned char> bytes(rawBytes, rawBytes + detailsData.length);
		NSDictionary* details = parseHorribleStructure(bytes);
		
		NSString* nsPreviewFileName = [[details objectForKey:NSN(3)] objectForKey:NSN(0)];
		NSArray* cacheFiles = [details objectForKey:NSN(10)];
		[crippled loadPreviewFromCacheFiles:cacheFiles previewName:nsPreviewFileName];
		
		return crippled.mapPreview;
	}
	
	return NULL;
}

-(NSString*)formattedVersion
{
	return [NSString stringWithFormat:@"%u.%u.%u.%u", version.major, version.minor, version.revision, version.build];
}

+(void)initialize
{
	NSMutableDictionary* parsers = [[NSMutableDictionary alloc] init];
	[parsers setObject:@"parseDetails:" forKey:kRPReplayDetailsKey];
	[parsers setObject:@"parseAttributeEvents:" forKey:kRPAttributeEventsKey];
	[parsers setObject:@"parseReplayInitData:" forKey:kRPReplayInitDataKey];
	[parsers setObject:@"parseGameEvents:" forKey:kRPReplayGameEventsKey];
	[parsers setObject:@"parseChatLog:" forKey:kRPReplayMessageEventsKey];
	
	fileParsers = [parsers copy];
	[parsers release];
}

-(id)crippledInitWithContentsOfFile:(NSString *)file
{
	if (!(self = [super init])) return nil;
	
	if (!SFileOpenArchive([file UTF8String], 0, MPQ_OPEN_READ_ONLY, &mpq))
	{
		[self release];
		return nil;
	}
	
	return self;
}

-(id)initWithContentsOfFile:(NSString *)file
{
	if (!(self = [super init])) return nil;
	
	if (!SFileOpenArchive([file UTF8String], 0, MPQ_OPEN_READ_ONLY, &mpq))
	{
		[self release];
		return nil;
	}
	
	// StormLib ignores and discards the user data header, so we need to retrieve it ourselves
	NSData* rawMpq = [[NSData alloc] initWithContentsOfMappedFile:file];
	const unsigned char* rawBytes = (const unsigned char*)rawMpq.bytes;
	checked_ptr<const unsigned char> bytes(rawBytes, rawBytes + rawMpq.length);
	bytes += 16;
	NSDictionary* headerInfo = parseHorribleStructure(bytes);
	[rawMpq release];
	
	NSDictionary* versionInfo = [headerInfo objectForKey:NSN(1)];
	version.major = [[versionInfo objectForKey:NSN(1)] unsignedIntValue];
	version.minor = [[versionInfo objectForKey:NSN(2)] unsignedIntValue];
	version.revision = [[versionInfo objectForKey:NSN(3)] unsignedIntValue];
	version.build = [[versionInfo objectForKey:NSN(4)] unsignedIntValue];
	
	duration = [[headerInfo objectForKey:NSN(3)] unsignedIntValue] / 16.0;
	
	for (NSString* key in fileParsers)
	{
		if (NSData* fileData = [self readArchivedFile:key])
		{
			SEL selector = NSSelectorFromString([fileParsers objectForKey:key]);
			[self performSelector:selector withObject:fileData];
		}
	}
	
	return self;
}

-(void)parseDetails:(NSData *)file
{
	const unsigned char* rawBytes = static_cast<const unsigned char*>(file.bytes);
	checked_ptr<const unsigned char> bytes(rawBytes, rawBytes + file.length);
	NSDictionary* details = parseHorribleStructure(bytes);
	
	// players
	NSArray* registeredPlayers = [details objectForKey:NSN(0)];
	NSMutableArray* playersBuilder = [NSMutableArray arrayWithCapacity:registeredPlayers.count];
	for (NSDictionary* registeredPlayer in registeredPlayers)
	{
		NSDictionary* identification = [registeredPlayer objectForKey:NSN(1)];
		unsigned uid = [[identification objectForKey:NSN(4)] intValue];
		RSPlayerRegion region = (RSPlayerRegion)[[identification objectForKey:NSN(0)] unsignedIntValue];
		NSString* name = [registeredPlayer objectForKey:NSN(0)];
		
		NSDictionary* color = [registeredPlayer objectForKey:NSN(3)];
		unsigned char r = [[color objectForKey:NSN(1)] unsignedCharValue];
		unsigned char g = [[color objectForKey:NSN(2)] unsignedCharValue];
		unsigned char b = [[color objectForKey:NSN(3)] unsignedCharValue];
		
		BOOL isWinner = [[registeredPlayer objectForKey:NSN(8)] intValue] == 1;
		
		RSPlayerData* player = [[RSPlayerData alloc] init];
		player.uid = uid;
		player.name = name;
		player.color = RSReplayColorCreate(r, g, b);
		player.isWinner = isWinner;
		player.region = region;
		
		[playersBuilder addObject:player];
		[player release];
	}
	
	players = [playersBuilder copy];
	
	// various metadata
	mapName = [[details objectForKey:NSN(1)] retain];
	endDate = [dateFromSC2Timestamp([[details objectForKey:NSN(5)] unsignedLongLongValue]) retain];
	startDate = [[NSDate alloc] initWithTimeInterval:-duration sinceDate:endDate];
	
	NSString* nsPreviewFileName = [[details objectForKey:NSN(3)] objectForKey:NSN(0)];
	NSArray* cacheFiles = [details objectForKey:NSN(10)];
	[self loadPreviewFromCacheFiles:cacheFiles previewName:nsPreviewFileName];
}

-(void)parseAttributeEvents:(NSData*)file
{
	NSMutableArray* teamConfigurations = [NSMutableArray arrayWithCapacity:7];
	for (int i = 0; i < 7; i++)
		[teamConfigurations addObject:[NSMutableArray array]];
	
	if (file.length < sizeof(RSReplayAttributesEventHeader))
	{
		RPLog(@"Malformed attributes stream");
		return;
	}
	
	const unsigned char* bytes = static_cast<const unsigned char*>(file.bytes);
	if (version.build > 17325) // for some reason Blizzard added a padding byte on build 17326
		bytes++;
	const RSReplayAttributesEventHeader* events = reinterpret_cast<const RSReplayAttributesEventHeader*>(bytes);
	
	NSUInteger eventStreamSize = sizeof (RSReplayAttributesEventHeader) + sizeof (RSReplayAttributesData) * events->attributesCount;
	if (file.length < eventStreamSize)
	{
		RPLog(@"Malformed attributes stream");
		return;
	}
	
	for (unsigned i = 0; i < events->attributesCount; i++)
	{
		const RSReplayAttributesData* event = events->attributes + i;
		RSPlayerData* player = event->playerId != 0x10 ? [players objectAtIndex:event->playerId - 1] : nil;
		switch (event->attributeId)
		{
			// player attributes
			case 0x01F4: // player type
				break; // we already know if the player is a comp
				
			case 0x0BB9: // player race
				switch (event->value)
				{
					case 'RAND': player.race = kRPPlayerRandomRace; break;
					case 'Terr': player.race = kRPPlayerTerranRace; break;
					case 'Prot': player.race = kRPPlayerProtossRace; break;
					case 'Zerg': player.race = kRPPlayerZergRace; break;
				}
				break;
				
			case 0x0BBA: // player color, already known
				break;
				
			case 0xBBB: // handicap
				player.handicap = event->value / 100.0;
				break;
				
			case 0x0BBC: // difficulty
				player.difficulty = (RSPlayerDifficulty)event->value;
				break;
				
			case 0x0BBF: // 'Part'
				if (event->value != 'Part')
				{
					RPLog(@"0BBF seen with value %X!", event->value);
				}
				break;
				
			case 0x0BC0: // 'Obs'
				if (event->value != 'Obs')
				{
					RPLog(@"0BC0 seen with value %04X!", event->value);
				}
				break;
				
			// team attributes
			case 0x07D2:
			case 0x07D3:
			case 0x07D4:
			case 0x07D5:
			case 0x07D6:
			case 0x07D7:
			case 0x07D8:
			{
				NSMutableArray* teamsBuilder = [teamConfigurations objectAtIndex:event->attributeId - 0x07D2];
				int team = (event->value & 0xff) - '1';
				if (team == -1) // everyone on a new team
					[teamsBuilder addObject:[NSArray arrayWithObject:player]];
				else
				{
					while (teamsBuilder.count <= team)
						[teamsBuilder addObject:[NSMutableArray array]];
					[[teamsBuilder objectAtIndex:team] addObject:player];
				}
				break;
			}
				
			case 0x7DB:
			case 0x7DC:
			case 0x7DD:
			case 0x7DE:
			case 0x7DF:
			case 0x7E0:
			case 0x7E1:
			case 0x7E2:
			case 0x7E3:
			case 0x7E4:
			case 0x7E5:
			case 0x7E6:
				break;
				
			// game attributes
			case 0x07D1: // game type
				type = (RSGameType)event->value;
				break;
				
			case 0x0BB8: // game speed
				speed = (RSGameSpeed)event->value;
				break;
				
			case 0x0BC1: // game ranking
				ranking = (RSGameRanking)event->value;
				break;
				
			case 0x0BC2:
				if (event->value != 'yes')
				{
					RPLog(@"%04X seen with value %04X!", event->attributeId, event->value);
				}
				break;
				
			case 0x03E8:
				if (event->value != 'Dflt')
				{
					RPLog(@"%04X seen with value %04X!", event->attributeId, event->value);
				}
				break;
				
			case 0x07D0:
				if (event->value != 't2')
				{
					RPLog(@"%04X seen with value %04X!", event->attributeId, event->value);
				}
				break;
				
			case 0x0BBE:
				if (event->value != '10')
				{
					RPLog(@"%04X seen with value %04X!", event->attributeId, event->value);
				}
				break;
				
			case 0x03E9: // either 'yes' or 'no'
				break;
				
			default:
				RPLog(@"Unknown attribute %04hx (player %i, value %08x)", (unsigned short)event->attributeId, event->playerId, event->value);
				break;
		}
	}
	
	assert(self.type != 0 && "No appropriate team configuration");
	NSMutableArray* teamsBuilder = nil;
	switch (self.type)
	{
		case RSGameType1v1: teamsBuilder = [teamConfigurations objectAtIndex:0]; break;
		case RSGameType2v2: teamsBuilder = [teamConfigurations objectAtIndex:1]; break;
		case RSGameType3v3: teamsBuilder = [teamConfigurations objectAtIndex:2]; break;
		case RSGameType4v4: teamsBuilder = [teamConfigurations objectAtIndex:3]; break;
		case RSGameTypeFFA: teamsBuilder = [teamConfigurations objectAtIndex:4]; break;
		case RSGameTypeCustom: teamsBuilder = [teamConfigurations objectAtIndex:5]; break;
		default: assert(!"confused about the game type"); break;
	}
	
	for (NSUInteger i = 0; i < teamsBuilder.count; i++)
	{
		NSMutableArray* team = [teamsBuilder objectAtIndex:i];
		[teamsBuilder replaceObjectAtIndex:i withObject:[[team copy] autorelease]];
	}
	teams = [teamsBuilder copy];
}

-(void)parseReplayInitData:(NSData*)file
{
	//RPLog(@"%s was called", sel_getName(_cmd));
}

-(void)parseGameEvents:(NSData*)file
{
	//RPLog(@"%s was called", sel_getName(_cmd));
}

-(void)parseChatLog:(NSData*)file
{
	const unsigned char* rawBytes = static_cast<const unsigned char*>(file.bytes);
	checked_ptr<const unsigned char> bytes(rawBytes, rawBytes + file.length);
	
	NSMutableArray* messagesBuilder = [NSMutableArray array];
	unsigned totalTime = 0;
	while (!bytes.ended())
	{
		unsigned timeDelta = parseTimeDelta(bytes);
		totalTime += timeDelta;
		
		int playerId = *bytes & 0xF;
		bytes++;
		
		unsigned char opcode = *bytes;
		bytes++;
		if ((opcode & 0x80) == 0)
		{
			unsigned target = opcode & 3;
			unsigned length = *bytes;
			bytes++;
			if (opcode & 8) length += 64;
			if (opcode & 16) length += 128;
			
			bytes.check(length);
			NSString* message = [[[NSString alloc] initWithBytes:bytes.to_unchecked() length:length encoding:NSUTF8StringEncoding] autorelease];
			bytes += length;
			
			// FIXME this assumes that player IDs are always consecutive
			RSPlayerData* player = [players objectAtIndex:playerId - 1];
			NSTimeInterval time = totalTime / 16;
			RSChatMessage* chatEntry = [[RSChatMessage alloc] initWithPlayer:player target:(RPReplayMessageTarget)target time:time message:message];
			[messagesBuilder addObject:chatEntry];
			[chatEntry release];
		}
		else if (opcode == 0x83)
		{
			// ping?
			bytes += 8;
		}
		else if (opcode == 0x80)
		{
			// weird header thing?
			bytes += 4;
		}
	}
	
	messages = [messagesBuilder copy];
}

-(NSData*)readArchivedFile:(NSString *)file
{
	void* fileHandle;
	const char* fileKey = [file UTF8String];
	if (!SFileHasFile(mpq, fileKey))
		return nil;
	
	if (SFileOpenFileEx(mpq, fileKey, SFILE_OPEN_FROM_MPQ, &fileHandle))
	{
		unsigned highSize;
		unsigned lowSize = SFileGetFileSize(fileHandle, &highSize);
		unsigned long long size = (static_cast<unsigned long long>(highSize) << 32) | lowSize;
		
		assert(size < UINT_MAX && "File is too large to be read");
		
		if (size == 0)
		{
			RPLog(@"File %@ is empty", file);
			SFileCloseFile(fileHandle);
			return nil;
		}
		
		unsigned char* buffer = static_cast<unsigned char*>(malloc(size));
		if (!SFileReadFile(fileHandle, buffer, static_cast<unsigned>(size)))
		{
			RPLog(@"Couldn't read %@: %s", file, GetMacOSStatusCommentString(GetLastError()));
			free(buffer);
			[self release];
			return nil;
		}
		
		NSData* result = [NSData dataWithBytes:buffer length:size];
		SFileCloseFile(fileHandle);
		free(buffer);
		return result;
	}
	else
	{
		RPLog(@"Couldn't open %@: %s", file, GetMacOSStatusCommentString(GetLastError()));
		return nil;
	}
}

-(void)loadPreviewFromCacheFiles:(NSArray*)cacheFiles previewName:(NSString *)previewFilePath
{
	const char* previewFileName = [previewFilePath UTF8String];
	for (NSData* cacheFileHash in cacheFiles)
	{
		NSData* actualHash = [cacheFileHash subdataWithRange:NSMakeRange(8, 32)];
		NSString* stringHash = stringnifyMapHash(actualHash);
		
		const char* bytes = (const char*)actualHash.bytes;
		NSString* resourcePath = [NSString stringWithFormat:kRSMapDataPathFormat, bytes[0], bytes[1], stringHash];
		
		void* mapData;
		if (!SFileOpenArchive([resourcePath UTF8String], 0, MPQ_OPEN_READ_ONLY, &mapData))
			continue;
		
		if (SFileHasFile(mapData, previewFileName))
		{
			// preview image found
			void* previewFile;
			if (SFileOpenFileEx(mapData, previewFileName, 0, &previewFile))
			{
				unsigned highSize;
				unsigned lowSize = SFileGetFileSize(previewFile, &highSize);
				if (highSize == 0) // we don't handle >4GB preview images, thank you
				{
					void* previewFileData = malloc(lowSize);
					if (!SFileReadFile(previewFile, previewFileData, lowSize))
					{
						RPLog(@"Couldn't read %@: %s", previewFilePath, GetMacOSStatusCommentString(GetLastError()));
					}
					
					// this data is either TGA or DDS; not very appealing in either case
					NSString* extension = [[previewFilePath pathExtension] lowercaseString];
					NSData* nsPreviewFileData = [NSData dataWithBytes:previewFileData length:lowSize];
					free(previewFileData);
					if ([extension isEqualToString:@"dds"])
					{
						mapPreview = CGImageRetain(loadDDS(nsPreviewFileData));
					}
					else if ([extension isEqualToString:@"tga"])
					{
						mapPreview = CGImageRetain(loadTGA(nsPreviewFileData));
					}
					else
					{
						RPLog(@"Unsupported preview type %@", extension);
					}
				}
				SFileCloseFile(previewFile);
			}
			else
			{
				RPLog(@"Couldn't open %@: %s", previewFilePath, GetMacOSStatusCommentString(GetLastError()));
			}
		}
		SFileCloseArchive(mapData);
		
		if (mapPreview != nil) break;
	}
}

-(void)dealloc
{
	SFileCloseArchive(mpq);
	CGImageRelease(mapPreview);
	[mapName release];
	[players release];
	[teams release];
	[messages release];
	[startDate release];
	[endDate release];
	[super dealloc];
}

@end
